#include "stdafx.h"
#include "FixedMemMgr.h"

#define PTRSIZE sizeof(void*)
#define ALIGNED_SIZE(Size, AlignTo) (Size + (AlignTo - ((Size - 1) & (AlignTo - 1))) - 1)
#define AlignedSize(Size) ALIGNED_SIZE(Size, PTRSIZE)
#define AlignedClassSize(C) AlignedSize(sizeof(C))
#define BlockBufferOffset AlignedClassSize(CBlock)
#define BlockBuffer(pBlock) ((BYTE*)pBlock + BlockBufferOffset)

HRESULT CFixedMemMgr::CreateInstance(const long ElementSize, const long ElementsPerBlock, FixedSizeMemoryManager** ppMemMgr)
{
    if (ElementSize < 1 || ElementsPerBlock < 1)
    {
        *ppMemMgr = NULL;
        return E_INVALIDARG;
    }

    BYTE* pNew;
    // Adjust to guarantee a reasonable alignment
    long AdjustedElemSize = AlignedSize(ElementSize);
    pNew = (BYTE*) new BYTE[AlignedClassSize(CFixedMemMgr) + PTRSIZE + AdjustedElemSize * ElementsPerBlock];
    if (NULL == pNew)
    {
        *ppMemMgr = NULL;
        return E_OUTOFMEMORY;
    }
    
    new ((void*)pNew) CFixedMemMgr();
    *ppMemMgr = (FixedSizeMemoryManager*)pNew;

    CFixedMemMgr& MemMgr = **((CFixedMemMgr**)ppMemMgr);
    MemMgr.m_ElemsPerBlock = MemMgr.m_cFreshItems = ElementsPerBlock;
    MemMgr.m_ElemSize = AdjustedElemSize;
    *(void**)(MemMgr.m_pFirstBlock = (void*)(pNew + AlignedClassSize(CFixedMemMgr))) = NULL;
    MemMgr.m_pFreshItem = (void*)((BYTE*)MemMgr.m_pFirstBlock + PTRSIZE);
    MemMgr.AddRef();
    return NOERROR;
}
CFixedMemMgr::~CFixedMemMgr()
{
    void* pBlock;
    void* pBlockNext = m_pFirstBlock;
    while (pBlock = pBlockNext)
    {
        pBlockNext = *(void**)pBlock;
        if (pBlockNext)
        {
            // Last block allocated with item, don't free it directly.
            delete pBlock;
        }
    }
}
bool CFixedMemMgr::AllocBlock()
{
    BYTE* pNew = new BYTE[PTRSIZE + m_ElemSize * m_ElemsPerBlock];
    if (pNew)
    {
        // First PTRSIZE bytes are the next pointer
        *(void**)pNew = (void*)m_pFirstBlock;
        m_pFirstBlock = (void*)pNew;

        // Point the list of fresh items at the first
        // slot in the new block.
        m_pFreshItem = (void*)(pNew + PTRSIZE);
        m_cFreshItems = m_ElemsPerBlock;
        return true;
    }
    return false;
}

STDMETHODIMP
CFixedMemMgr::Alloc(long *retVal)
{
    // See if we have any items that have been freed.
    // Give them right back out if we do.
    if (m_pFree)
    {
        *retVal = (long)m_pFree;
        m_pFree = *(void**)m_pFree;
    }
    else
    {
        if (!m_cFreshItems && !AllocBlock())
        {
            return E_OUTOFMEMORY;
        }
        *retVal = (long)m_pFreshItem;
        m_pFreshItem = (void*)((BYTE*)m_pFreshItem + m_ElemSize);
        m_cFreshItems--;
    }

    // Zero out our allocated memory
    ZeroMemory((void*)*retVal, m_ElemSize);

    return NOERROR;
}
STDMETHODIMP_(void)
CFixedMemMgr::Free(long pMemory)
{
    // Just chain back in
    // Note, no check is done on pMemory in this system
    *(void**)pMemory = m_pFree;
    m_pFree = (void*)pMemory;
}
STDMETHODIMP
CFixedMemMgr::get_ElementSize(long *retVal)
{
    *retVal = m_ElemSize;
    return NOERROR;
}
STDMETHODIMP
CFixedMemMgr::get_ElementsPerBlock(long *retVal)
{
    *retVal = m_ElemsPerBlock;
    return NOERROR;
}

HRESULT CFixedMemMgrCompact::CreateInstance(const long ElementSize, const long ElementsPerBlock, FixedSizeMemoryManager** ppMemMgr)
{
    HRESULT hr;
    if (ElementSize < 1 || ElementsPerBlock < 1)
    {
        *ppMemMgr = NULL;
        return E_INVALIDARG;
    }
    if (SUCCEEDED(hr = (*ppMemMgr = new CFixedMemMgrCompact()) ? NOERROR : E_OUTOFMEMORY))
    {
        CFixedMemMgrCompact& MemMgr = **((CFixedMemMgrCompact**)ppMemMgr);
        MemMgr.m_ElemsPerBlock = ElementsPerBlock;
        // Adjust to get a reasonable alignment
        MemMgr.m_ElemSize = AlignedSize(ElementSize);
        MemMgr.m_BlockSize = BlockBufferOffset + MemMgr.m_ElemSize * ElementsPerBlock;
        MemMgr.AddRef();
    }
    return hr;
}
CFixedMemMgrCompact::~CFixedMemMgrCompact()
{
    CBlock* pBlock;
    CBlock* pBlockNext = m_pFirstBlock;
    while (pBlock = pBlockNext)
    {
        pBlockNext = pBlock->pNext;
        delete (void*)pBlock;
    }
}
CFixedMemMgrCompact::CBlock* CFixedMemMgrCompact::AllocBlock()
{
    CBlock* pNewBlock = (CBlock*) new BYTE[m_BlockSize];
    if (pNewBlock)
    {
        CBlock& Block = *pNewBlock;
        Block.pNext = m_pFirstBlock;
        m_pFirstBlock = pNewBlock;
        Block.cInUse = 0;
        Block.fFreeIsList = false;
        Block.pFree = (void*)BlockBuffer(pNewBlock);
    }
    return pNewBlock;
}

STDMETHODIMP
CFixedMemMgrCompact::Alloc(long *retVal)
{
    // Find the first block with an available slot
    CBlock* pBlock = m_pFirstBlock;
    CBlock* pBlockPrev = NULL;
    while (pBlock && !pBlock->pFree)
    {
        pBlockPrev = pBlock;
        pBlock = pBlock->pNext;
    }

    // Allocate a new block if one isn't available
    if (pBlock)
    {
        if (!pBlock->cInUse)
        {
            m_cEmptyBlocks--;
        }
    }
    else
    {
        if (NULL == (pBlock = AllocBlock()))
        {
            *retVal = NULL;
            return E_OUTOFMEMORY;
        }
        pBlockPrev = NULL;
    }

    // Move this block to the head of the list
    // so the next search is more efficient.
    if (pBlockPrev)
    {
        pBlockPrev->pNext = pBlock->pNext;
        pBlock->pNext = m_pFirstBlock;
        m_pFirstBlock = pBlock;
    }

    // Get a piece of memory from the block.  This is
    // guaranteed to return valid memory at this point.
    CBlock& Block = *pBlock;
    *retVal = (long)Block.pFree;

    // Update the free list for the block.
    if (++Block.cInUse == m_ElemsPerBlock)
    {
        // All elements have been allocated, none are free.
        Block.pFree = NULL;
    }
    else if (Block.fFreeIsList)
    {
        // Grab the next item from the free list
        if (NULL == (Block.pFree = *(void**)Block.pFree))
        {
            // Revert back to in order allocation if there are no
            // more items in the list.
            Block.fFreeIsList = false;
            Block.pFree = (void*)(BlockBuffer(pBlock) +  Block.cInUse * m_ElemSize);
        }
    }
    else
    {
        // We're allocating in order
        Block.pFree = (void*)((BYTE*)Block.pFree + m_ElemSize);
    }

    // Zero out our allocated memory
    ZeroMemory((void*)*retVal, m_ElemSize);

    return NOERROR;
}
STDMETHODIMP_(void)
CFixedMemMgrCompact::Free(long pMemory)
{
    CBlock* pBlock = m_pFirstBlock;
    CBlock* pBlockPrev = NULL;
    while (pBlock)
    {
        if ((BYTE*)pMemory > (BYTE*)pBlock && (BYTE*)pMemory < ((BYTE*)pBlock + m_BlockSize))
        {
            break;
        }
        pBlockPrev = pBlock;
        pBlock = pBlock->pNext;
    }

    // If we found a block, then we're in business
    if (pBlock)
    {
        // If this isn't the first block, then move it to the front of the
        // list just like in the alloc case.  Allocs and Frees tend to happen
        // in batches, and this will speed things significantly in the normal case.
        if (pBlockPrev)
        {
            pBlockPrev->pNext = pBlock->pNext;
            pBlock->pNext = m_pFirstBlock;
            m_pFirstBlock = pBlock;
        }

        // Free the item
        *(void**)pMemory = pBlock->fFreeIsList ? pBlock->pFree : NULL;
        pBlock->fFreeIsList = true;
        pBlock->pFree = (void*)pMemory;
        if (0 == --pBlock->cInUse)
        {
            // This block is no longer in use.  See if we should
            // free it right away.
            if (m_fCompactOnFree && m_cEmptyBlocks >= m_cBufferBlocks)
            {
                // We've already moved this block to the front, so relinking the
                // list is trivial.
                m_pFirstBlock = pBlock->pNext;
                delete (void*)pBlock; 
            }
            else
            {
                m_cEmptyBlocks++;
            }
        }
    }
    // Note: Silent failure.  Free doesn't return an HRESULT for call speed
    // reasons, and CFixedMemMgr::Free can't return one anyway.
}
STDMETHODIMP
CFixedMemMgrCompact::get_ElementSize(long *retVal)
{
    *retVal = m_ElemSize;
    return NOERROR;
}
STDMETHODIMP
CFixedMemMgrCompact::get_ElementsPerBlock(long *retVal)
{
    *retVal = m_ElemsPerBlock;
    return NOERROR;
}
STDMETHODIMP_(void)
CFixedMemMgrCompact::Compact(void)
{
    int i = m_cEmptyBlocks - m_cBufferBlocks;
    if (i > 0)
    {
        CBlock* pBlock = m_pFirstBlock;
        CBlock** ppThisBlock = &m_pFirstBlock;
        m_cEmptyBlocks -= (short)i;
        // while (pBlock) is unnecessary here, the loop will exit before pBlock is NULL
        for (;;)
        {
            if (!pBlock->cInUse)
            {
                *ppThisBlock = pBlock->pNext;
                delete (void*)pBlock;
                pBlock = *ppThisBlock;
                if (0 == --i)
                {
                    break;
                }
            }
            else
            {
                ppThisBlock = &pBlock->pNext;
                pBlock = pBlock->pNext;
            }
        }
    }
}
STDMETHODIMP
CFixedMemMgrCompact::get_CompactOnFree(VARIANT_BOOL *retVal)
{
    *retVal = m_fCompactOnFree;
    return NOERROR;
}
STDMETHODIMP
CFixedMemMgrCompact::put_CompactOnFree(VARIANT_BOOL RHS)
{
    m_fCompactOnFree = RHS;
    return NOERROR;
}
STDMETHODIMP
CFixedMemMgrCompact::get_BufferBlocks(short *retVal)
{
    *retVal = m_cBufferBlocks;
    return NOERROR;
}
STDMETHODIMP
CFixedMemMgrCompact::put_BufferBlocks(short RHS)
{
    if (RHS < 0)
    {
        return E_INVALIDARG;
    }
    m_cBufferBlocks = RHS;
    return NOERROR;
}
